import json
import os
import socket
from fl.logger import logger

class SocketServer:
    """
    SocketServer类用于创建一个TCP服务器，用于接收客户端的连接和数据传输。
    """
    def __init__(self, host:str, port:int, upload_directory:str, download_directory:str):
        self.host = host  # 监听所有可用的接口
        self.port = port  # 监听的端口
        self.upload_directory = upload_directory
        self.download_directory = download_directory
        self.server_socket = None

    def ensure_directory_exists(self):
        if not os.path.exists(self.upload_directory):
            os.makedirs(self.upload_directory)

    def start_server(self):
        self.server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.server_socket.bind((self.host, self.port))
        self.server_socket.listen(5)
        logger.info(f"Listening on {self.host}:{self.port}")

        while True:
            client_socket, addr = self.server_socket.accept()
            self.handle_client(client_socket)
            logger.info(f"Connection from {addr}")

    def handle_client(self, client_socket):
        try:
            # 接收请求类型长度和请求类型
            request_type_length = int.from_bytes(self.receive_data(client_socket, 4), 'big')
            request_type = self.receive_data(client_socket, request_type_length).decode('utf-8')
            logger.info(f"Received request type: {request_type}")

            if request_type == 'upload':
                # 接收文件名长度和文件名
                filename_length = int.from_bytes(self.receive_data(client_socket, 4), 'big')
                filename = self.receive_data(client_socket, filename_length).decode('utf-8')
                logger.info(f"Received file name length: {filename_length}")

                # 接收文件内容长度和文件内容
                file_content_length = int.from_bytes(self.receive_data(client_socket, 4), 'big')
                file_content = self.receive_data(client_socket, file_content_length)
                logger.info(f"Received file content length: {file_content_length}")

                # 接收传感器数据长度和传感器数据
                # sensor_data_length = int.from_bytes(self.receive_data(client_socket, 4), 'big')
                # logger.info(f"Received sensor data length: {sensor_data_length}")
                # sensor_data = self.receive_data(client_socket, sensor_data_length).decode('utf-8')

                # 确保目录存在
                self.ensure_directory_exists()

                # 保存文件
                file_path = os.path.join(self.upload_directory, filename)
                with open(file_path, 'wb') as file:
                    file.write(file_content)

                # sensor_data_path = os.path.join(self.upload_directory, filename + '.json')
                # with open(sensor_data_path, 'w') as sensor_file:
                #     json.dump(json.loads(sensor_data), sensor_file)

                client_socket.send(b'{"status": "ok", "message": "File successfully uploaded"}')
            elif request_type == 'download':
                # 接收文件名长度和文件名
                filename_length = int.from_bytes(self.receive_data(client_socket, 4), 'big')
                filename = self.receive_data(client_socket, filename_length).decode('utf-8')

                file_path = os.path.join(self.download_directory, filename)
                if not os.path.exists(file_path):
                    client_socket.send(b'{"status": "error", "message": "File not found"}')
                    return

                # 读取文件内容
                with open(file_path, 'rb') as file:
                    file_content = file.read()

                # 发送文件内容长度和文件内容
                client_socket.send(len(file_content).to_bytes(4, 'big'))
                client_socket.send(file_content)

                client_socket.send(b'{"status": "ok", "message": "File successfully downloaded"}')
            else:
                client_socket.send(b'{"status": "error", "message": "Invalid request type"}')
                logger.error("Invalid request type")
        except Exception as e:
            client_socket.send(b'{"status": "error", "message": "Error occurred"}')
            logger.error(f"Error occurred while uploading file: {e}")
        finally:
            client_socket.close()

    def receive_data(self, sock, length):
        data = bytearray()
        while len(data) < length:
            packet = sock.recv(min(length - len(data), 4096))  # 每次接收最多 4096 字节
            if not packet:
                raise Exception("Connection lost while receiving data")
            data.extend(packet)
        return data




